import re
import asyncio
import websockets
from datetime import datetime


CONNECTIONS = set()


async def send(websocket, message):
    try:
        await websocket.send(message)
    except websockets.ConnectionClosed:
        pass


# 广播任务, 即向所有的websocket链接发送新用户加入的消息
async def broadcast():
    for websocket in CONNECTIONS:
        await asyncio.wait([
            asyncio.create_task(send(websocket, f"{datetime.now()}: New user `{websocket.host}` joined"))
        ])
    await asyncio.sleep(3)

async def handler(websocket):
    loop = asyncio.get_running_loop()
    loop.create_task(broadcast())           # 添加新的异步广播任务

    data = await websocket.recv()
    CONNECTIONS.add(websocket)              # 这里添加websocket到链接池
    print("Receive:", data)
    try:
        client = re.findall(re.compile(r"I am\s+([a-z,A-Z]+)"), data)[0]
    except IndexError:
        client = "?"
    reply = f"{datetime.now()}: Hello {client}, I'm luffy, "
    await websocket.send(reply)
    print(f"Response:{reply}")

    try:
        await websocket.wait_closed()
    finally:
        CONNECTIONS.remove(websocket)


async def main():
    async with websockets.serve(handler, "localhost", 9999):
        await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())
